#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <math.h>

typedef struct body body;
typedef struct node node;

#define SIZE 9999.0
#define THETA_SQUARE 0.25
#define SOFTEN 18
#define DAMPING 0.1

double time_step;
unsigned int depth;

struct body {
  unsigned int num;

  unsigned int mass;

  double loc_x;
  double loc_y; 

  double vel_x;
  double vel_y;
  double accln_x;
  double accln_y;
};

struct node {
  node* children[4];
  node* parent;
  body* b[4];

  unsigned int mass;

  double x_start;
  double x_end;
  double y_start;
  double y_end;

  double loc_x;
  double loc_y;
};

void* populate_list(FILE* ifile, int* count) {
  unsigned int i;
  fscanf(ifile, "%u", count);

  body* list = malloc(sizeof(body)*(*count));

  for(i=0; i<*count; i++) {
    list[i].num = i;
    fscanf(ifile, "%lf %lf %u %*u %*u", &list[i].loc_x, &list[i].loc_y, &list[i].mass);
    list[i].vel_x = 0;
    list[i].vel_y = 0;
    list[i].accln_x = 0;
    list[i].accln_y = 0;
  }

  return list;
}

void inline init_node(node* n, node* p, double x_start, double x_end, double y_start, double y_end) {
  n->children[0] = NULL;  n->children[1] = NULL;
  n->children[2] = NULL;  n->children[3] = NULL;
  n->parent = p;
  n->b[0] = NULL;  n->b[1] = NULL;
  n->b[2] = NULL;  n->b[3] = NULL;
  n->x_start = x_start;
  n->x_end = x_end;
  n->y_start = y_start;
  n->y_end = y_end;
}

void build_tree_helper(node* n, body* b) {
  double x = b->loc_x;
  double y = b->loc_y;
  double x_mid = (n->x_start + n->x_end)/2;
  double y_mid = (n->y_start + n->y_end)/2;

  if (depth > 100)
    return;
  depth++;

  if (x_mid - x > 0) {
    if (y_mid - y > 0) {
      // quadrature 1
      if (n->b[0] == NULL) {
        if (n->children[0] == NULL) {
          n->b[0] = b;
          return;
        }
        else {
          return build_tree_helper(n->children[0], b);
        }
      }
      else {
        n->children[0] = malloc(sizeof(node));
        init_node(n->children[0], n, n->x_start, x_mid, n->y_start, y_mid);
        build_tree_helper(n->children[0], n->b[0]);
        n->b[0] = NULL;
        return build_tree_helper(n->children[0], b);
      }
    }
    // quadrature 2
    else {
      if (n->b[1] == NULL) {
        if (n->children[1] == NULL) {
          n->b[1] = b;
          return;
        }
        else {
          return build_tree_helper(n->children[1], b);
        }
      }
      else {
        n->children[1] = malloc(sizeof(node));
        init_node(n->children[1], n, n->x_start, x_mid, y_mid, n->y_end);
        build_tree_helper(n->children[1], n->b[1]);
        n->b[1] = NULL;
        return build_tree_helper(n->children[1], b);
      }
    }
  }

  else {
    // quadrature 3
    if (y_mid - y > 0) {
      if (n->b[2] == NULL) {
        if (n->children[2] == NULL) {
          n->b[2] = b;
          return;
        }
        else {
          return build_tree_helper(n->children[2], b);
        }
      }
      else {
        n->children[2] = malloc(sizeof(node));
        init_node(n->children[2], n, x_mid, n->x_end, n->y_start, y_mid);
        build_tree_helper(n->children[2], n->b[2]);
        n->b[2] = NULL;
        return build_tree_helper(n->children[2], b);
      }
    }
    // quadrature 4
    else {
      if (n->b[3] == NULL) {
        if (n->children[3] == NULL) {
          n->b[3] = b;
          return;
        }
        else {
          return build_tree_helper(n->children[3], b);
        }
      }
      else {
        n->children[3] = malloc(sizeof(node));
        init_node(n->children[3], n, x_mid, n->x_end, y_mid, n->y_end);
        build_tree_helper(n->children[3], n->b[3]);
        n->b[3] = NULL;
        return build_tree_helper(n->children[3], b);
      }
    }
  }
}

void inline build_tree(node* root, body* list, int count) {
  int i=0;
  for(i=0; i<count; i++) {
    depth = 0;
    build_tree_helper(root, &(list[i]));
  }
}

void sweep_up(node* root) {
  int i;
  root->mass = 0;

  for(i=0; i<4; i++) {
    if (root->children[i] != NULL) {
      sweep_up(root->children[i]);
      root->loc_x = root->mass * root->loc_x + root->children[i]->mass * root->children[i]->loc_x;
      root->loc_y = root->mass * root->loc_y + root->children[i]->mass * root->children[i]->loc_y;
      root->mass += root->children[i]->mass;
    }
    else if (root->b[i] != NULL) {
      root->loc_x = root->mass * root->loc_x + root->b[i]->mass * root->b[i]->loc_x;
      root->loc_y = root->mass * root->loc_y + root->b[i]->mass * root->b[i]->loc_y;
      root->mass += root->b[i]->mass;
    }
  }
}

void cal_force_helper(node* n, body* b, unsigned int size) {
  double mag;
  double x_diff = (n->loc_x - b->loc_x);
  double y_diff = (n->loc_y - b->loc_y);
  double distance = x_diff*x_diff + y_diff*y_diff;

  double size_squre = size*size;

  if (distance/size_squre < THETA_SQUARE) {
    // use approximation
    distance += SOFTEN;
    distance = pow(distance, 1.5);
    mag = n->mass / distance;

    b->accln_x += mag * x_diff;
    b->accln_y += mag * y_diff;
    return;
  }

  else {
    int i=0;
    for(i; i<4; i++) {
      if (n->b[i] != NULL) {
        x_diff = (n->b[i]->loc_x - b->loc_x);
        y_diff = (n->b[i]->loc_y - b->loc_y);
        distance = x_diff*x_diff + y_diff*y_diff;
        distance += SOFTEN;
        distance = pow(distance, 1.5);
        mag = n->b[i]->mass / distance;

        b->accln_x += mag * x_diff;
        b->accln_y += mag * y_diff;
      }

      if (n->children[i] != NULL) {
        cal_force_helper(n->children[i], b, size/2);
      }
    }
  }

}

void cal_force(node* root, body* list, unsigned int count) {
  int i;

  for(i=0; i<count; i++) {
    list[i].accln_x = 0;
    list[i].accln_y = 0;
    cal_force_helper(root, &(list[i]), SIZE);
    printf("%lf ", list[i].accln_y);
  }
  printf("\n");

}

void time_advance(node* root) {
  int i=0;
  for(i; i<4; i++) {
    if (root->b[i] != NULL) {
      root->b[i]->vel_x += root->b[i]->accln_x*time_step;
      root->b[i]->vel_x *= DAMPING;
      root->b[i]->vel_y += root->b[i]->accln_y*time_step;
      root->b[i]->vel_y *= DAMPING;

      root->b[i]->loc_x += root->b[i]->vel_x * time_step;
      root->b[i]->loc_y += root->b[i]->vel_y * time_step;

      if (root->b[i]->loc_x > SIZE)
        root->b[i]->loc_x = SIZE;
      if (root->b[i]->loc_y > SIZE)
        root->b[i]->loc_y = SIZE;
      if (root->b[i]->loc_x < -SIZE)
        root->b[i]->loc_x = -SIZE;
      if (root->b[i]->loc_x < -SIZE)
        root->b[i]->loc_x = -SIZE;
    }

    if (root->children[i] != NULL)
      time_advance(root->children[i]);
  }
  free(root);
  return;
}

int main(int argc, char* argv[]) {
  assert(argc == 5 && "need extra input parameters");

  double num_ite = atof(argv[1]);
  time_step = atof(argv[2]);
  double time = 0;
  FILE* ifile = fopen(argv[3], "r");
  body* list;
  unsigned int count, i=0;

  // read info about all input bodies
  list = populate_list(ifile, &count);
  fclose (ifile);


  while(i < num_ite) {
    // create root for the tree
    node* root = malloc(sizeof(node));
    init_node(root, NULL, -SIZE, SIZE, -SIZE, SIZE);

    // build the tree using body list
    build_tree(root, list, count);

    // calculate center loc and mass bottom up
    sweep_up(root);

    // calculate force for each body in the tree
    cal_force(root, list, count);

    // time advance
    time_advance(root);

    i++;
  }

  FILE* ofile = fopen(argv[4], "w");
  fprintf(ofile, "%u\n", count);
  for(i=0; i<count; i++) {
    fprintf(ofile, "%.2lf %.2lf %u %.2lf %.2lf\n", list[i].loc_x, list[i].loc_y, list[i].mass, list[i].vel_x, list[i].vel_y);
  }
  fclose(ofile);

  
  return 0;
}
